# -*- coding: utf-8 -*-
"""
按 DOY（年内天序）逐点学习 2016–2021 的年际趋势 → 外推得到 2022 的等距预测（每4天）
加入：
  - 趋势缩放 TREND_DAMPING: y22 = y21 + λ*(y22_raw - y21)
  - 鲁棒回归（Theil-Sen/RANSAC），减弱异常年份影响
流程：训练年(2016–2021)窗口回归后插值 → 组训练矩阵(行=2022网格DOY,列=年份)
    → 逐 DOY 拟合 Year→Value → 外推2022并做趋势缩放 → Savitzky–Golay平滑
    → 在 test.xlsx 的日期上用 PCHIP 取值 → 评估并导出

输出：
- output_p1/train_matrix_every4d.csv
- output_p1/pred_2022_grid_every4d.csv（含 raw/smoothed/damped 三列）
- output_p1/pred_eval_2022_by_DOY_damped.csv
"""

# =============== 参数区（可改） ===============
INPUT_XLSX = "P1sortdata.xlsx"   # 全量原始数据
SHEET_NAME = 0

TEST_XLSX = "test.xlsx"          # 2022 测试集（两列：日期, 值[可为0~1或百分比]）
TEST_YEAR = 2022

# 窗口与分辨率
WINDOW_START = (3, 1)            # 仅建模 3/1～7/1
WINDOW_END   = (7, 1)
RESOLUTION_DAYS = 4              # 统一等间隔（每4天）

# 每年在窗口内的“回归后插值”方式： "poly" 或 "loess"
REGRESSION_KIND = "poly"
POLY_DEGREE_CANDIDATES = (1, 2, 3)
LOESS_FRAC = 0.6

# 逐 DOY 的“年际趋势”模型：
YEAR_TREND = "linear"            # "linear" | "quadratic" | "isotonic"
ROBUST = "theilsen"              # "none" | "theilsen" | "ransac"  (仅对 linear 生效)
RIDGE_ALPHA = 0.0                # 线性/二次的 L2 正则（>0 生效）

# 趋势缩放：λ∈(0,1]，越小越保守，避免外推过冲
TREND_DAMPING = 0.05

# 对 2022 等距预测序列的平滑
SMOOTH = True
SMOOTH_WINDOW = 9                # 奇数；过大自动调整
SMOOTH_POLY   = 2

# 结果合理性保护（可选）：整体带宽裁剪（均值±k*std）
USE_GLOBAL_BAND = True
GLOBAL_BAND_SIGMA = 2.5

OUTPUT_DIR = "./output_p1"
# ============================================

import os, warnings
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d, PchipInterpolator
from scipy.signal import savgol_filter
from sklearn.linear_model import LinearRegression, Ridge, TheilSenRegressor, RANSACRegressor
from sklearn.isotonic import IsotonicRegression
from sklearn.preprocessing import PolynomialFeatures
from sklearn.pipeline import Pipeline
warnings.filterwarnings("ignore")
os.makedirs(OUTPUT_DIR, exist_ok=True)

# ---------------- I/O 与预处理 ----------------
def _infer_cols(df: pd.DataFrame):
    date_candidates  = ["日期","date","时间","Date","time","Time"]
    value_candidates = ["就业率","就业率(%)","就业率（%）","y","value","rate"]
    dcol = None
    for c in df.columns:
        try:
            pd.to_datetime(df[c], errors="raise"); dcol = c; break
        except Exception: pass
    if dcol is None:
        for c in date_candidates:
            if c in df.columns: dcol = c; break
    vcol = None
    for c in df.columns:
        if c == dcol: continue
        s = pd.to_numeric(df[c], errors="coerce")
        if s.notna().sum() >= max(5, int(0.3*len(s))):
            vcol = c
            if any(k in str(c) for k in ["就","rate","value","%"]): break
    if vcol is None:
        for c in value_candidates:
            if c in df.columns: vcol = c; break
    if dcol is None or vcol is None:
        raise ValueError(f"无法识别日期/数值列：{df.columns.tolist()}")
    return dcol, vcol

def load_all(path, sheet=0):
    df = pd.read_excel(path, sheet_name=sheet)
    dcol, vcol = _infer_cols(df)
    df = df[[dcol, vcol]].copy()
    df.columns = ["date","y"]
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["y"]    = pd.to_numeric(df["y"], errors="coerce")
    df = df.dropna().drop_duplicates("date").sort_values("date").reset_index(drop=True)
    # 单位统一到百分比
    if df["y"].median() <= 1.5 and df["y"].max() <= 1.2:
        df["y"] *= 100.0
    df["y"] = df["y"].clip(0, 100)
    df["year"] = df["date"].dt.year
    return df

def load_test(path):
    df = pd.read_excel(path)
    dcol, vcol = _infer_cols(df)
    df = df[[dcol, vcol]].copy()
    df.columns = ["date","y_true"]
    df["date"] = pd.to_datetime(df["date"], errors="coerce")
    df["y_true"] = pd.to_numeric(df["y_true"], errors="coerce")
    df = df.dropna().sort_values("date").reset_index(drop=True)
    if df["y_true"].median() <= 1.5 and df["y_true"].max() <= 1.2:
        df["y_true"] *= 100.0
    df["y_true"] = df["y_true"].clip(0, 100)
    return df

def make_grid(year, start_mmdd, end_mmdd, step_days):
    start = pd.Timestamp(year=year, month=start_mmdd[0], day=start_mmdd[1])
    end   = pd.Timestamp(year=year, month=end_mmdd[0],   day=end_mmdd[1])
    return pd.date_range(start, end, freq=f"{int(step_days)}D")

# ---------------- 每年回归后插值（统一到等间隔） ----------------
try:
    from statsmodels.nonparametric.smoothers_lowess import lowess
except Exception:
    lowess = None

def _aicc(n, sse, k):
    if n <= k + 1: return np.inf
    return n*np.log(max(sse/n, 1e-12)) + 2*k + (2*k*(k+1))/(n - k - 1)

def fit_poly_reg(x, y, deg, alpha=1e-6):
    mdl = Pipeline([
        ("poly", PolynomialFeatures(degree=deg, include_bias=True)),
        ("ridge", Ridge(alpha=alpha, fit_intercept=False, random_state=0))
    ])
    mdl.fit(x.reshape(-1,1), y)
    yhat = mdl.predict(x.reshape(-1,1))
    sse = float(np.sum((y - yhat)**2))
    return mdl, _aicc(len(y), sse, deg + 1)

def per_year_regression_interpolate(df_all, step_days, start_mmdd, end_mmdd, kind="poly"):
    years = sorted(df_all["year"].unique())
    target_grid = make_grid(TEST_YEAR, start_mmdd, end_mmdd, step_days)
    per_year_series = {}
    for y in years:
        g = df_all[df_all["year"] == y].copy()
        if g.empty: continue
        start = pd.Timestamp(year=y, month=start_mmdd[0], day=start_mmdd[1])
        end   = pd.Timestamp(year=y, month=end_mmdd[0],   day=end_mmdd[1])
        g_win = g[(g["date"] >= start) & (g["date"] <= end)].copy()
        g_use = g_win if len(g_win) >= 2 else g.copy()
        grid = make_grid(y, start_mmdd, end_mmdd, step_days)
        if len(g_use) < 2:
            if len(g) >= 2:
                xi = g["date"].astype("int64").values
                yi = g["y"].values
                try:   f = interp1d(xi, yi, kind="linear",  fill_value="extrapolate")
                except Exception: f = interp1d(xi, yi, kind="nearest", fill_value="extrapolate")
                yg = f(grid.astype("int64").values)
            else:
                yg = np.repeat(g["y"].iloc[0], len(grid))
            per_year_series[y] = pd.Series(np.clip(yg,0,100), index=grid); continue

        x = g_use["date"].dt.dayofyear.values.astype(float)
        yv = g_use["y"].values.astype(float)
        xg = pd.DatetimeIndex(grid).dayofyear.values.astype(float)
        if kind == "loess" and lowess is not None and len(g_use) >= 5:
            sm = lowess(yv, x, frac=LOESS_FRAC, it=0, return_sorted=True)
            f = interp1d(sm[:,0], sm[:,1], kind="linear", fill_value="extrapolate")
            yg = f(xg)
        else:
            best = (None, np.inf, None)
            for deg in POLY_DEGREE_CANDIDATES:
                try:
                    m, aicc = fit_poly_reg(x, yv, deg)
                    if aicc < best[1]: best = (deg, aicc, m)
                except Exception:
                    continue
            model = best[2] if best[2] is not None else Ridge().fit(x.reshape(-1,1), yv)
            yg = model.predict(xg.reshape(-1,1))
        per_year_series[y] = pd.Series(np.clip(yg,0,100), index=grid)

    # 以 2022 的 grid（DOY）为基准拼矩阵
    rows = []
    for dt in target_grid:
        doy = dt.timetuple().tm_yday
        row = {"date": dt}
        for y, ser in per_year_series.items():
            ser_doy = ser.index.dayofyear.values
            pos = np.where(ser_doy == doy)[0]
            j = int(pos[0]) if len(pos) else int(np.argmin(np.abs(ser_doy - doy)))
            row[y] = float(ser.iloc[j])
        rows.append(row)
    mat_df = pd.DataFrame(rows).set_index("date").sort_index()
    return mat_df

# ---------------- 逐 DOY 年际趋势拟合（含鲁棒&缩放） ----------------
def fit_year_trend(x_years, y_values, year_trend="linear", robust="none", ridge_alpha=0.0):
    x = np.asarray(x_years, dtype=float).reshape(-1,1)
    y = np.asarray(y_values, dtype=float)
    mask = np.isfinite(y)
    if mask.sum() < 2:
        return None  # 不足以拟合
    x, y = x[mask], y[mask]

    if year_trend == "linear":
        if robust == "theilsen":
            mdl = TheilSenRegressor(random_state=0)
            mdl.fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))
        elif robust == "ransac":
            base = LinearRegression()
            mdl = RANSACRegressor(base_estimator=base, random_state=0)
            mdl.fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))
        else:
            if ridge_alpha and ridge_alpha > 0:
                mdl = Ridge(alpha=ridge_alpha, fit_intercept=True, random_state=0).fit(x, y)
            else:
                mdl = LinearRegression().fit(x, y)
            return lambda yr: float(mdl.predict(np.array([[yr]], dtype=float)))

    elif year_trend == "quadratic":
        X = np.hstack([x, x**2])
        if ridge_alpha and ridge_alpha > 0:
            mdl = Ridge(alpha=ridge_alpha, fit_intercept=True, random_state=0).fit(X, y)
        else:
            mdl = LinearRegression().fit(X, y)
        return lambda yr: float(mdl.predict(np.array([[yr, yr**2]], dtype=float)))

    elif year_trend == "isotonic":
        # 年份单调同增
        order = np.argsort(x[:,0])
        iso = IsotonicRegression(increasing=True, out_of_bounds="clip")
        iso.fit(x[order,0], y[order])
        return lambda yr: float(iso.predict([yr])[0])

    else:
        raise ValueError("YEAR_TREND 必须是 'linear' | 'quadratic' | 'isotonic'")

def smooth_series(y, window=9, poly=2):
    n = len(y)
    if n < 5: return y
    w = min(window, n - (1 - n % 2))
    if w < 3: w = 3
    if w % 2 == 0: w -= 1
    if poly >= w: poly = max(1, w - 1)
    try:
        return savgol_filter(y, window_length=w, polyorder=poly, mode="interp")
    except Exception:
        return y

def apply_global_band(y, k=2.5):
    mu = np.nanmean(y); sd = np.nanstd(y)
    if not np.isfinite(sd) or sd == 0: return y
    lo, hi = mu - k*sd, mu + k*sd
    return np.clip(y, lo, hi)

def metrics(y_true, y_pred):
    y_true = np.asarray(y_true); y_pred = np.asarray(y_pred)
    mse  = np.mean((y_pred - y_true)**2)
    rmse = np.sqrt(mse)
    mae  = np.mean(np.abs(y_pred - y_true))
    eps = 1e-6
    mape = np.mean(np.abs((y_pred - y_true) / np.maximum(np.abs(y_true), eps))) * 100.0
    smape = np.mean(2*np.abs(y_pred - y_true) / np.maximum(np.abs(y_pred)+np.abs(y_true), eps)) * 100.0
    ss_res = np.sum((y_pred - y_true)**2); ss_tot = np.sum((y_true - y_true.mean())**2)
    r2 = 1 - ss_res/ss_tot if ss_tot > 0 else np.nan
    return {"MSE": mse, "RMSE": rmse, "MAE": mae, "MAPE(%)": mape, "sMAPE(%)": smape, "R2": r2}

# ---------------- 主流程 ----------------
def main():
    # 1) 读数据
    df_all = load_all(INPUT_XLSX, SHEET_NAME)
    test_df = load_test(TEST_XLSX)

    # 2) 训练年：2016–2021
    train_all = df_all[(df_all["year"] >= 2016) & (df_all["year"] < TEST_YEAR)].copy()
    if train_all.empty: raise RuntimeError("训练集为空。")

    # 3) 每年窗口内回归后插值 → 训练矩阵（行=2022的DOY网格，列=年份）
    mat_df = per_year_regression_interpolate(
        train_all, RESOLUTION_DAYS, WINDOW_START, WINDOW_END, kind=REGRESSION_KIND
    )
    mat_df.to_csv(os.path.join(OUTPUT_DIR, f"train_matrix_every{RESOLUTION_DAYS}d.csv"),
                  encoding="utf-8-sig")

    years = np.array(sorted([c for c in mat_df.columns if isinstance(c, (int,np.integer))]), dtype=float)
    grid2022 = mat_df.index

    # 4) 逐 DOY 拟合 Year→Value，并做趋势缩放
    y22_raw = []
    y22_damped = []
    for _, row in mat_df.iterrows():
        vals = row[years].values.astype(float)
        pred_fn = fit_year_trend(years, vals, year_trend=YEAR_TREND, robust=ROBUST, ridge_alpha=RIDGE_ALPHA)
        if pred_fn is None:
            # 兜底：用非空均值
            m = float(np.nanmean(vals)) if np.isfinite(np.nanmean(vals)) else 0.0
            y22_raw.append(m); y22_damped.append(m); continue
        y21 = float(pred_fn(2021))
        y22r = float(pred_fn(2022))
        y22r = np.clip(y22r, 0, 100)
        # 趋势缩放
        y22adj = y21 + TREND_DAMPING * (y22r - y21)
        y22adj = np.clip(y22adj, 0, 100)
        y22_raw.append(y22r); y22_damped.append(y22adj)

    y22_raw = np.array(y22_raw, dtype=float)
    y22_damped = np.array(y22_damped, dtype=float)

    # 5) 全局带宽保护（可选）
    if USE_GLOBAL_BAND:
        y22_damped = apply_global_band(y22_damped, k=GLOBAL_BAND_SIGMA)
        y22_damped = np.clip(y22_damped, 0, 100)

    # 6) 平滑
    y22_sm = y22_damped.copy()
    if SMOOTH:
        y22_sm = smooth_series(y22_sm, window=SMOOTH_WINDOW, poly=SMOOTH_POLY)
        y22_sm = np.clip(y22_sm, 0, 100)

    # 7) 保存 2022 等距预测（原始外推、缩放后、平滑后）
    pd.DataFrame({
        "date": grid2022,
        "yhat_2022_raw(%)": y22_raw,
        "yhat_2022_damped(%)": y22_damped,
        "yhat_2022_smooth(%)": y22_sm
    }).to_csv(os.path.join(OUTPUT_DIR, f"pred_2022_grid_every{RESOLUTION_DAYS}d.csv"),
              index=False, encoding="utf-8-sig")

    # 8) 在 test.xlsx 的日期上用 PCHIP 取值并评估（基于平滑后的序列）
    st = pd.Timestamp(year=TEST_YEAR, month=WINDOW_START[0], day=WINDOW_START[1])
    ed = pd.Timestamp(year=TEST_YEAR, month=WINDOW_END[0],   day=WINDOW_END[1])
    test_df = test_df[(test_df["date"] >= st) & (test_df["date"] <= ed)].copy()
    if test_df.empty: raise RuntimeError("测试点不在 3/1–7/1 窗口内。")

    x_grid = grid2022.view("int64").astype(float)
    x_test = test_df["date"].view("int64").astype(float)
    try:
        f = PchipInterpolator(x_grid, y22_sm, extrapolate=True)
        y_pred = f(x_test)
    except Exception:
        f = interp1d(x_grid, y22_sm, kind="linear", fill_value="extrapolate")
        y_pred = f(x_test)

    y_true = test_df["y_true"].values
    score = metrics(y_true, y_pred)

    out = pd.DataFrame({
        "date": test_df["date"],
        "y_true(%)": y_true,
        "y_pred(%)": y_pred,
        "abs_err": np.abs(y_pred - y_true)
    })
    out.to_csv(os.path.join(OUTPUT_DIR, "pred_eval_2022_by_DOY_damped.csv"),
               index=False, encoding="utf-8-sig")

    print("\n=== Evaluation on 2022 (Per-DOY Year-Trend + Damping) ===")
    for k, v in score.items():
        print(f"{k}: {v:.6f}")
    print("\n关键参数：")
    print(f" - YEAR_TREND={YEAR_TREND}, ROBUST={ROBUST}, TREND_DAMPING={TREND_DAMPING}")
    print(f" - SMOOTH={SMOOTH}, WINDOW={SMOOTH_WINDOW}, POLY={SMOOTH_POLY}")
    print("\n输出文件：")
    print(f" - 训练矩阵：{os.path.join(OUTPUT_DIR, f'train_matrix_every{RESOLUTION_DAYS}d.csv')}")
    print(f" - 2022等距预测：{os.path.join(OUTPUT_DIR, f'pred_2022_grid_every{RESOLUTION_DAYS}d.csv')}")
    print(f" - 测试评估：{os.path.join(OUTPUT_DIR, 'pred_eval_2022_by_DOY_damped.csv')}")
    print("\n小贴士：若3月初仍偏高，优先把 TREND_DAMPING 调小（如 0.3），或将 ROBUST 设为 'ransac'。")

if __name__ == "__main__":
    main()
